import numpy
import math
import warnings
import lxml.etree as ET

# constants
slotSize = 16.0
hullCenter = 256.0
maxDimension = 1024.0
slotRestrictions = ("I", "O", "E")
"""
slotRestrictionSets = []
for i in range(2**len(slotRestrictions)):
    s = ""
    for j, slotRestriction in enumerate(slotRestrictions):
        if i & (1 << j): s += slotRestriction
    slotRestrictionSets.append(s)
slotRestrictionSets = tuple(slotRestrictionSets)
"""

slotToModelRotation = numpy.array([
    [0.0, 1.0, 0.0],
    [-1.0, 0.0, 0.0],
    [0.0, 0.0, 1.0]])

def slotsFromCenter(moduleSlotData):
    position = moduleSlotData.find("Position")
    slotX = ( float(position.findtext("X")) - hullCenter ) / slotSize
    slotY = ( float(position.findtext("Y")) - hullCenter ) / slotSize
    return slotX, slotY

def slotsFromCorner(moduleSlotData, nx, ny):
    position = moduleSlotData.find("Position")
    ix = (float(position.findtext("X")) - hullCenter - slotSize/2.0) / slotSize + nx/2.0
    iy = (float(position.findtext("Y")) - hullCenter - slotSize/2.0) / slotSize + ny/2.0
    return ix, iy

def hullDimensions(moduleSlotDatas):
    # returns the bounding box of the ship and the size of the ship
    # in terms of slots
    # slots: a list of <ModuleSlotData> Elements
    sizeX = 0.0
    sizeY = 0.0
    for moduleSlotData in moduleSlotDatas:
        slotX, slotY = slotsFromCenter(moduleSlotData)
        distX = 2.0 * abs(slotX) + 1.0
        distY = 2.0 * abs(slotY) + 1.0
        sizeX = max(distX, sizeX)
        sizeY = max(distY, sizeY)
    return math.ceil(sizeX), math.ceil(sizeY)

def slots2ImageData(slots, channelOrder=slotRestrictions):
    nx, ny = slots.shape
    result = numpy.zeros((nx, ny, 4), int)

    slotRestrictionIndices = []
    for channel in channelOrder:
        if channel in slotRestrictions:
            slotRestrictionIndices.append(slotRestrictions.index(channel))
        else:
            slotRestrictionIndices.append(None)
            
    for i, sri in enumerate(slotRestrictionIndices):
        if sri is None: continue
        compareSlots = numpy.zeros_like(slots)
        compareSlots.fill(1 << sri)
        match = numpy.bitwise_and(slots, compareSlots) > 0
        
        result[match, i] = 255
        result[match, 3] = 255
    return result

def string2Slot(s):
    result = 0
    for i, slotChar in enumerate(slotRestrictions):
        if slotChar in s: result |= (1 << i)
    return result

def shipData2Slots(shipData):
    moduleSlotList = shipData.find("ModuleSlotList")
    moduleSlotDatas = moduleSlotList.findall("ModuleSlotData")

    sizeX, sizeY = hullDimensions(moduleSlotDatas)
    result = numpy.zeros((sizeX, sizeY), int)

    for moduleSlotData in moduleSlotDatas:
        slotString = moduleSlotData.findtext("Restrictions").upper()
        slot = string2Slot(slotString)
        ix, iy = slotsFromCorner(moduleSlotData, sizeX, sizeY)
        if ix % 1.0 or iy % 1.0:
            warnMsg = 'Misaligned slot would be placed at array position (%f, %f); truncating to integer (%.0f, %.0f).' % (ix, iy, ix, iy)
            warnings.warn(warnMsg, stacklevel=2)
            ix = round(ix)
            iy = round(iy)
        result[ix, iy] = slot
    return result

def imageData2Slots(imageData, channelOrder=slotRestrictions):
    # converts an image to slot types
    nx, ny, nc = imageData.shape
    result = numpy.zeros((nx, ny), int)
    
    slotRestrictionIndices = []
    for channel in channelOrder:
        if channel in slotRestrictions:
            slotRestrictionIndices.append(slotRestrictions.index(channel))
        else:
            slotRestrictionIndices.append(None)
    
    for i, sri in enumerate(slotRestrictionIndices):
        if sri is None: continue
        result[imageData[:, :, i] > 127] += (1 << sri)
    return result

def imageDataIndex2HullCoordinate(ix, iy, nx, ny):
    hx = (ix + 0.5 - 0.5 * nx) * slotSize + hullCenter
    hy = (iy + 0.5 - 0.5 * ny) * slotSize + hullCenter
    return hx, hy

def slot2String(slotType):
    result = ""
    for i, slotChar in enumerate(slotRestrictions):
        if slotType & (1 << i): result += slotChar
    return result

def countSlotRestrictionSets(slots):
    # counts the number of slots with a particular set of restrictions
    # returns empty, I, O, IO, E, IE, OE, IOE
    n = 2**len(slotRestrictions)
    result = tuple(numpy.count_nonzero(slots == i) for i in range(n))
    return result

def countSlotRestrictions(slots):
    # counts the number of slots with a particular restriction
    # returns total slots, (number of I, number of O, number of E)
    # an IO slot would be counted as both I and O
    result = [0 for x in slotRestrictions]
    numTotal = numpy.count_nonzero(slots)
    for i, slotChar in enumerate(slotRestrictions):
        compareSlots = numpy.zeros_like(slots)
        compareSlots.fill(1 << i)
        result[i] = numpy.count_nonzero(numpy.bitwise_and(slots, compareSlots))
    return numTotal, tuple(result)
        
def hullRole(slots):
    # returns a hull type given a set of slots
    numTotal, (numI, numO, numE) = countSlotRestrictions(slots)
    # "mobile" ships have at least 1% engines
    if float(numO) / numTotal < 0.01:
        # emplacements
        if numTotal < 400:
            return "platform"
        else:
            return "station"
    else:
        # freighter: less than 40% O slots
        if float(numO) / numTotal < 0.4:
            return "freighter"
        
        # combat hulls
        if numTotal < 50:
            return "fighter"
        if numTotal < 150:
            return "frigate"
        if numTotal < 400:
            return "cruiser"

        # arbitrary distinction between carrier and capital
        if numO > numI:
            return "carrier"
        else:
            return "capital"

def slots2ModuleSlotList(slots):
    nx, ny = slots.shape
    moduleSlotList = ET.Element("ModuleSlotList")
    for ix in range(nx):
        for iy in range(ny):
            slotString = slot2String(slots[ix, iy])
            if slotString == "": continue

            hx, hy = imageDataIndex2HullCoordinate(ix, iy, nx, ny)
            
            moduleSlotData = ET.SubElement(moduleSlotList, "ModuleSlotData")
            
            position = ET.SubElement(moduleSlotData, "Position")
            ET.SubElement(position, "X").text = "%d" % hx
            ET.SubElement(position, "Y").text = "%d" % hy

            ET.SubElement(moduleSlotData, "Health").text = "0"
            ET.SubElement(moduleSlotData, "Shield_Power").text = "0"
            ET.SubElement(moduleSlotData, "facing").text = "0"
            ET.SubElement(moduleSlotData, "Restrictions").text = slotString
    return moduleSlotList

def defaultHullXML(slots, name):
    nsmap = {
        "xsi" : "http://www.w3.org/2001/XMLSchema-instance",
        "xsd" : "http://www.w3.org/2001/XMLSchema",
        }
    shipData = ET.Element("ShipData", nsmap = nsmap)
    ET.SubElement(shipData, "Name").text = name
    ET.SubElement(shipData, "Hull").text = name
    ET.SubElement(shipData, "Role").text = hullRole(slots)
    ET.SubElement(shipData, "IconPath").text = 'ShipIcons/%s' % name
    ET.SubElement(shipData, "ThrusterList")
    ET.SubElement(shipData, "ModelPath").text = 'Model/Ships/%s' % name
    ET.SubElement(shipData, "DefaultAIState").text = "AwaitingOrders"
    ET.SubElement(shipData, "CombatState").text = "Artillery"
    shipData.append(slots2ModuleSlotList(slots))
    
    return shipData

def countAsymmetricalSlotPairs(slots):
    nx, ny = slots.shape
    result = numpy.count_nonzero(slots[:nx//2, :] != slots[:-1-nx//2:-1, :])
    return result
